

import os 
import sys 


root='/share/project/cwm/shaocong.xu/exp/stable-normal-2'
os.chdir(root)


sys.path.append(root)
from daniel_tools.img_utils import *
from  os.path import dirname

from tqdm import tqdm

parser = argparse.ArgumentParser()

parser.add_argument('--img_bank', type=str, 
default='/share/project/cwm/shaocong.xu/exp/stable-normal-2/checkpoints/stabletransparent-512X512-cleargrasp-depth/image_log/train', help='the image resource ')


args = parser.parse_args()





# img_bank = '/share/project/cwm/shaocong.xu/exp/stable-normal-2/checkpoints/stabletransparent-512X512-cleargrasp-depth-sd-naive-from-sd2.1/image_log/train'
img_bank = args.img_bank

coarse_samples_gs = sorted([join(img_bank,x) for x in os.listdir(img_bank) if x.startswith('coarse_samples_gs')])
fine_samples_gs = sorted([ join(img_bank,x) for x in os.listdir(img_bank) if x.startswith('fine_samples_gs')])
control_gs = sorted([ join(img_bank,x) for x in os.listdir(img_bank) if x.startswith('control_gs')])
reconstruction_gs = sorted([ join(img_bank,x) for x in os.listdir(img_bank) if x.startswith('reconstruction_gs')])





merged_list = []

for  ref, gt,  coarse, fine in tqdm(zip(control_gs,reconstruction_gs,coarse_samples_gs,fine_samples_gs)):


    tmp = ref.split('/')[-1].split('.')[0].split('-')
    iter_num, epoch_num = tmp[1],tmp[2]

    ref_img = cv2.imread(ref)
    gt_img = cv2.imread(gt)
    coarse_img = cv2.imread(coarse)
    fine_img = cv2.imread(fine)

    # ref_img = cv2.resize(ref_img,(512,512))
    # gt_img = cv2.resize(gt_img,(512,512))
    # coarse_img = cv2.resize(coarse_img,(512,512))
    # fine_img = cv2.resize(fine_img,(512,512))

    ref_img = cv2.cvtColor(ref_img,cv2.COLOR_BGR2RGB)
    gt_img = cv2.cvtColor(gt_img,cv2.COLOR_BGR2RGB)
    coarse_img = cv2.cvtColor(coarse_img,cv2.COLOR_BGR2RGB)
    fine_img = cv2.cvtColor(fine_img,cv2.COLOR_BGR2RGB)

    merged = np.concatenate([ref_img,gt_img,coarse_img,fine_img],axis=0)




    text = f"{iter_num}-{epoch_num}"  # 要添加的文字
    position = (300,100)  # (x, y) 左下角起点
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 1.5
    color = (0, 0, 255)  # BGR格式（红色）
    thickness = 2

    
    cv2.putText(merged, text, position, font, font_scale, color, thickness)

    # cv2.imwrite(f'merged-{iter_num}-{epoch_num}.png',merged)

    merged_list.append(merged)




imgs2video2(merged_list, out_path=join(dirname(img_bank),'merged.mp4'))

# images_to_gif(fine_samples_gs,output_gif=join(dirname(img_bank),'fine_samples_gs.gif'))
# images_to_gif(fine_samples_gs,output_gif= join(dirname(img_bank),'coarse_samples_gs.gif'))


# imgs2video(fine_samples_gs,out_path=join(dirname(img_bank),'fine_samples_gs.mp4'))
# imgs2video(fine_samples_gs,out_path= join(dirname(img_bank),'coarse_samples_gs.mp4'))

